package com.gitee.feizns.explore.data_structure.union_find;

/**
 * @author feizns
 * @since 2020/3/11
 */
public class UnionFindSizeOptimization implements UnionFind {

    private int[] parent;
    private int[] sizes;

    public UnionFindSizeOptimization(int size) {
        this.parent = new int[size];
        this.sizes = new int[size];
        for (int i = 0; i < size; i++) {
            parent[i] = i;
            sizes[i] = 1;
        }
    }

    @Override
    public int size() {
        return parent.length;
    }

    @Override
    public boolean isConnected(int p, int q) {
        return find(p) == find(q);
    }

    @Override
    public void union(int p, int q) {
        int pr = find(p);
        int qr = find(q);
        if ( pr != qr ) {
            if ( sizes[pr] <= sizes[qr] ) {
                parent[pr] = qr;
                sizes[qr] += sizes[pr];
            } else {
                parent[qr] = pr;
                sizes[pr] += sizes[qr];
            }
        }
    }

    private int find(int p) {
        int r = p;
        while ( r != parent[r] )
            r = parent[r];
        return r;
    }

    public static void main(String[] args) {
        UnionFind ou = new UnionFindSizeOptimization(20);
        ou.union(1, 19);
        ou.union(0, 1);

        ou.union(19, 10);
        System.out.println(ou.isConnected(10, 19));
    }

}
